import warnings
warnings.filterwarnings('ignore')
import os
import cv2

import numpy as np
import torch
import glob
from model.edge_model import SegUNet
from torchvision import transforms
import time
from PIL import Image
from collections import OrderedDict
from matplotlib import pyplot as plt
import onnx
from onnxsim import simplify

#model = TWoStageSSD(num_class=num_class, fpn_nums=3,scores_the=0.4, top_detect=2000)
model = SegUNet(9, is_edge=True)
model.load_state_dict(torch.load('/data/check/edge_model_1219_544att_306.pth'),strict=True)
#model.load_state_dict(torch.load('/home/dsl/release/xag_fpv_1207_scale2267.pth'),strict=True)


model.eval()
img = torch.randn(1,3, 544, 960, requires_grad=True)
torch_out = model(img)
torch.onnx.export(model,               # model being run
                  img,           # model input (or a tuple for multiple inputs)
                  "seg.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=11,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  verbose=False,
                 )

onnx_model = onnx.load("seg.onnx")
onnx.checker.check_model(onnx_model)
onnx_model = onnx.load('seg.onnx')
sim, _ =simplify(onnx_model)
onnx.save_model(sim, 'seg.onnx')